import numpy as np
import matplotlib.pyplot as plt
from plyfile import PlyData, PlyElement
import tkinter as tk
from tkinter import ttk
import math
from tkinter import messagebox 

color_map = [
    (128, 0, 0),     (132, 0, 0),     (136, 0, 0),     (140, 0, 0),
    (144, 0, 0),     (148, 0, 0),     (152, 0, 0),     (156, 0, 0),
    (160, 0, 0),     (164, 0, 0),     (168, 0, 0),     (172, 0, 0),
    (176, 0, 0),     (180, 0, 0),     (184, 0, 0),     (188, 0, 0),
    (192, 0, 0),     (196, 0, 0),     (200, 0, 0),     (204, 0, 0),
    (208, 0, 0),     (212, 0, 0),     (216, 0, 0),     (220, 0, 0),
    (224, 0, 0),     (228, 0, 0),     (232, 0, 0),     (236, 0, 0),
    (240, 0, 0),     (244, 0, 0),     (248, 0, 0),     (252, 0, 0),
    (255, 0, 0),     (255, 4, 0),     (255, 8, 0),     (255, 12, 0),
    (255, 16, 0),    (255, 20, 0),    (255, 24, 0),    (255, 28, 0),
    (255, 32, 0),    (255, 36, 0),    (255, 40, 0),    (255, 44, 0),
    (255, 48, 0),    (255, 52, 0),    (255, 56, 0),    (255, 60, 0),
    (255, 64, 0),    (255, 68, 0),    (255, 72, 0),    (255, 76, 0),
    (255, 80, 0),    (255, 84, 0),    (255, 88, 0),    (255, 92, 0),
    (255, 96, 0),    (255, 100, 0),   (255, 104, 0),   (255, 108, 0),
    (255, 112, 0),   (255, 116, 0),   (255, 120, 0),   (255, 124, 0),
    (255, 128, 0),   (255, 132, 0),   (255, 136, 0),   (255, 140, 0),
    (255, 144, 0),   (255, 148, 0),   (255, 152, 0),   (255, 156, 0),
    (255, 160, 0),   (255, 164, 0),   (255, 168, 0),   (255, 172, 0),
    (255, 176, 0),   (255, 180, 0),   (255, 184, 0),   (255, 188, 0),
    (255, 192, 0),   (255, 196, 0),   (255, 200, 0),   (255, 204, 0),
    (255, 208, 0),   (255, 212, 0),   (255, 216, 0),   (255, 220, 0),
    (255, 224, 0),   (255, 228, 0),   (255, 232, 0),   (255, 236, 0),
    (255, 240, 0),   (255, 244, 0),   (255, 248, 0),   (255, 252, 0),
    (254, 255, 1),   (250, 255, 6),   (246, 255, 10),  (242, 255, 14),
    (238, 255, 18),  (234, 255, 22),  (230, 255, 26),  (226, 255, 30),
    (222, 255, 34),  (218, 255, 38),  (214, 255, 42),  (210, 255, 46),
    (206, 255, 50),  (202, 255, 54),  (198, 255, 58),  (194, 255, 62),
    (190, 255, 66),  (186, 255, 70),  (182, 255, 74),  (178, 255, 78),
    (174, 255, 82),  (170, 255, 86),  (166, 255, 90),  (162, 255, 94),
    (158, 255, 98),  (154, 255, 102), (150, 255, 106), (146, 255, 110),
    (142, 255, 114), (138, 255, 118), (134, 255, 122), (130, 255, 126),
    (126, 255, 130), (122, 255, 134), (118, 255, 138), (114, 255, 142),
    (110, 255, 146), (106, 255, 150), (102, 255, 154), (98, 255, 158),
    (94, 255, 162),  (90, 255, 166),  (86, 255, 170),  (82, 255, 174),
    (78, 255, 178),  (74, 255, 182),  (70, 255, 186),  (66, 255, 190),
    (62, 255, 194),  (58, 255, 198),  (54, 255, 202),  (50, 255, 206),
    (46, 255, 210),  (42, 255, 214),  (38, 255, 218),  (34, 255, 222),
    (30, 255, 226),  (26, 255, 230),  (22, 255, 234),  (18, 255, 238),
    (14, 255, 242),  (10, 255, 246),  (6, 255, 250),   (2, 255, 254),
    (0, 252, 255),   (0, 248, 255),   (0, 244, 255),   (0, 240, 255),
    (0, 236, 255),   (0, 232, 255),   (0, 228, 255),   (0, 224, 255),
    (0, 220, 255),   (0, 216, 255),   (0, 212, 255),   (0, 208, 255),
    (0, 204, 255),   (0, 200, 255),   (0, 196, 255),   (0, 192, 255),
    (0, 188, 255),   (0, 184, 255),   (0, 180, 255),   (0, 176, 255),
    (0, 172, 255),   (0, 168, 255),   (0, 164, 255),   (0, 160, 255),
    (0, 156, 255),   (0, 152, 255),   (0, 148, 255),   (0, 144, 255),
    (0, 140, 255),   (0, 136, 255),   (0, 132, 255),   (0, 128, 255),
    (0, 124, 255),   (0, 120, 255),   (0, 116, 255),   (0, 112, 255),
    (0, 108, 255),   (0, 104, 255),   (0, 100, 255),   (0, 96, 255),
    (0, 92, 255),    (0, 88, 255),    (0, 84, 255),    (0, 80, 255),
    (0, 76, 255),    (0, 72, 255),    (0, 68, 255),    (0, 64, 255),
    (0, 60, 255),    (0, 56, 255),    (0, 52, 255),    (0, 48, 255),
    (0, 44, 255),    (0, 40, 255),    (0, 36, 255),    (0, 32, 255),
    (0, 28, 255),    (0, 24, 255),    (0, 20, 255),    (0, 16, 255),
    (0, 12, 255),    (0, 8, 255),     (0, 4, 255),     (0, 0, 255),
    (0, 0, 252),     (0, 0, 248),     (0, 0, 244),     (0, 0, 240),
    (0, 0, 236),     (0, 0, 232),     (0, 0, 228),     (0, 0, 224),
    (0, 0, 220),     (0, 0, 216),     (0, 0, 212),     (0, 0, 208),
    (0, 0, 204),     (0, 0, 200),     (0, 0, 196),     (0, 0, 192),
    (0, 0, 188),     (0, 0, 184),     (0, 0, 180),     (0, 0, 176),
    (0, 0, 172),     (0, 0, 168),     (0, 0, 164),     (0, 0, 160),
    (0, 0, 156),     (0, 0, 152),     (0, 0, 148),     (0, 0, 144),
    (0, 0, 140),     (0, 0, 136),     (0, 0, 132),     (0, 0, 128)
]

def visualize_ply(file_path):
    # 读取PLY文件
    ply_data = PlyData.read(file_path)
    vertices = ply_data['vertex']
    
    # 提取坐标数据
    x = vertices['x']
    y = vertices['y']
    z = vertices['z']
    
    # 计算深度范围
    min_depth = np.min(z)
    max_depth = np.max(z)
    depth_range = max_depth - min_depth
    
    # 生成颜色映射
    colors = []
    for depth in z:
        # 将深度映射到颜色索引 (0-255)
        color_idx = int(((depth - min_depth) / depth_range) * (len(color_map)-1))
        colors.append(color_map[color_idx])
    
    # 转换为0-1范围的RGB值
    colors = np.array(colors) / 255.0

    # 创建3D可视化
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制点云
    ax.scatter(x, y, z, c=colors, s=1, marker='o')
    
    # 设置坐标轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    
    # 启用交互式旋转
    plt.show(block=True)  # 添加block=True参数保持窗口显示

# 在文件末尾添加以下新功能
def generate_cylinder(radius, fov_deg, max_range, points=5000,noise=5,coefficient=4):
    if radius <= 0 or fov_deg <= 0 or max_range <= 0:
        return None

    if fov_deg > 180:
        return None

    if radius >= max_range:
        return None

    if points <= 0:
        return None

    min_range = radius / math.sin(math.radians(fov_deg / 2))
    if min_range > max_range:
        return None
    
    max_distance = math.sqrt(max_range**2 - radius**2)
    min_distance = radius / math.tan(math.radians(fov_deg / 2))

    print(max_distance)
    if max_distance <= min_distance or max_distance <= 0 or min_distance <= 0:
        return None

    a_rad = math.asin(radius / max_range)
    a_deg = math.degrees(a_rad)
    if a_deg > fov_deg / 2:
        return None

    # 生成圆柱面点云
    theta = np.random.uniform(0, 2*math.pi, points)

    u = np.random.rand(points)
    z = max_distance - (max_distance - min_distance) * np.power(u,1/coefficient)
    # z = np.random.uniform(min_distance, max_distance, points)

    # 转换为笛卡尔坐标（管壁表面）
    radius_noise = np.random.normal(0, noise, points)
    radius += radius_noise
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)

    x_noise = np.random.normal(0, noise, points)
    y_noise = np.random.normal(0, noise, points) 

    x +=  x_noise
    y +=  y_noise

    return x, y, z

class CylinderGUI:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title('圆柱体内部3D点云生成器')
        
        ttk.Label(self.root, text='管道半径:').grid(row=0, column=0)
        self.radius_entry = ttk.Entry(self.root)
        self.radius_entry.insert(0, '300')
        self.radius_entry.grid(row=0, column=1)

        # 更新输入控件
        ttk.Label(self.root, text='FOV角度:').grid(row=1, column=0)
        self.fov_entry = ttk.Entry(self.root)
        self.fov_entry.insert(0, '60')
        self.fov_entry.grid(row=1, column=1)
        
        ttk.Label(self.root, text='测量范围:').grid(row=2, column=0)
        self.range_entry = ttk.Entry(self.root)
        self.range_entry.insert(0, '2000')
        self.range_entry.grid(row=2, column=1)
        
        # 生成按钮
        self.generate_btn = ttk.Button(self.root, text='生成', command=self.generate)
        self.generate_btn.grid(row=3, columnspan=2)

    def generate(self):
        try:
            fov = float(self.fov_entry.get())
            max_range = float(self.range_entry.get())
            radius = float(self.radius_entry.get())

            # 调用新的生成函数
            li = generate_cylinder(
                radius=radius,
                fov_deg=fov,
                max_range=max_range
            )

            if li is None:
                messagebox.showwarning("警告", "未生成有效点云，请调整检测参数")
                return
            
            x, y, z = li

            # 创建PLY结构并保存（保留原有实现）
            vertices = np.zeros(len(x), dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')])
            vertices['x'], vertices['y'], vertices['z'] = x, y, z
            
            ply = PlyData([PlyElement.describe(vertices, 'vertex')], text=True)
            ply.write('tube.ply')
            
            # 调用可视化
            visualize_ply('tube.ply')
            
        except ValueError as e:
            print(f"参数错误: {e}")

if __name__ == "__main__":
    gui = CylinderGUI()
    gui.root.mainloop()

